import numpy as np
import math



class Memory(object):
    def __init__(self, capacity, dims):
        self.capacity = capacity
        self.data = np.zeros((capacity, dims))
        self.pointer = 0

    def store_transition(self, s, a, r, s_):
        transition = np.hstack((s, a[0][0],a[1][0], [r], s_))
        index = self.pointer % self.capacity  # replace the old memory with new memory
        self.data[index, :] = transition
        self.pointer += 1

    def sample(self, n):
        assert self.pointer >= self.capacity, '存储尚未填满'
        indices = np.random.choice(self.capacity, size=n)
        return self.data[indices, :]

class preprocess():
    """
    认为相对位置更加重要，坐标就不再作为输入了
    """
    def __init__(self,observation):
        """
        当前为输入敌我双方位置和弹药量，之后要再加入导弹位置
        数据进行归一化
        :param observation:
        """
        self.location_F=observation[0]
        self.location_E=observation[1]
        self.n_ammo=observation[2]
        self.pos_missile_e=observation[5]
        self.pointing=observation[10]


    def get_input(self):
        dis=self.cal_distance()
        hight=self.cal_diff_hight()
        ang=self.cal_ang()
        pointing=self.cal_pointing()
        ammo=self.n_ammo[0]/4       #   TODO  剩余导弹个数的归一化  2
        result=np.array([dis,hight,ang,pointing,ammo])
        return result

    def cal_distance(self):
        """
        坐标算距离
        :param a:
        :param b:
        :return:
        """
        a=self.location_F[0]
        b=self.location_E[0]
        diff=(a[0]-b[0])**2+(a[1]-b[1])**2+(a[2]-b[2])**2
        result=math.sqrt(diff)
        return self.nomal_atan(result,20)

    def cal_diff_hight(self):
        a = self.location_F[0]
        b = self.location_E[0]
        c=a[2]-b[2]
        return self.nomal_line1(c,20,0)

    def cal_ang(self):
        """
        计算目标点的偏航角
        :return:
        """
        a = self.location_F[0]
        b = self.location_E[0]
        y=b[1]-a[1]
        x=b[0]-a[0]
        ang=math.atan2(y,x)/math.pi
        return ang

    def cal_pointing(self):
        a=self.pointing[0][1]//math.pi
        return a/math.pi/2

    def cal_mdistance(self):
        """
        坐标算距离
        :param a:
        :param b:
        :return:
        """
        a=self.location_F[0]
        b=self.location_E[0]
        diff=(a[0]-b[0])**2+(a[1]-b[1])**2+(a[2]-b[2])**2
        result=math.sqrt(diff)
        return result

    def cal_mdiff_hight(self):
        a = self.location_F[0]
        b = self.location_E[0]
        return a[2]-b[2]

    def cal_mang(self):
        """
        计算目标点的偏角
        :return:
        """
        a = self.location_F[0]
        b = self.location_E[0]
        y=b[1]-a[1]
        x=b[0]-a[0]
        ang=math.atan2(y,x)
        return ang

    def nomal_log(self,data,i):
        return math.log(data+1,i)

    def nomal_atan(self,data,x_range):
        b=x_range/4
        d=data/b-2
        result=(math.atan(d)*2)/math.pi
        return result

    def nomal_line1(self,data,dmax,dmin):
        return (data-dmin)/(dmax-dmin)


